{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "672d679d",
   "metadata": {},
   "source": [
    "# 注册算子属性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1d437369",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.ir import Op\n",
    "import tvm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f01aec3c",
   "metadata": {},
   "source": [
    "## 测试算子属性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "85a1e488",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_op = Op.get(\"relax.log\")\n",
    "\n",
    "@tvm.ir.register_op_attr(\"exp\", \"ftest\")\n",
    "def test(x):\n",
    "    return x + 1\n",
    "\n",
    "assert log_op.num_inputs == 1\n",
    "assert log_op.get_attr(\"ftest\") is None\n",
    "assert Op.get(\"exp\").get_attr(\"ftest\")(1) == 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ade8b95",
   "metadata": {},
   "source": [
    "## 重置算子属性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d434640d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add1(x):\n",
    "    return x + 1\n",
    "\n",
    "def add2(x):\n",
    "    return x + 2\n",
    "\n",
    "# Register fadd1 and fadd2 attributes.\n",
    "tvm.ir.register_op_attr(\"exp\", \"fadd1\", add1)\n",
    "tvm.ir.register_op_attr(\"log\", \"fadd1\", add1)\n",
    "tvm.ir.register_op_attr(\"log\", \"fadd2\", add2)\n",
    "\n",
    "# Reset log fadd1 attr.\n",
    "log_op = Op.get(\"log\")\n",
    "log_op.reset_attr(\"fadd1\")\n",
    "\n",
    "# Check that fadd1 attr is resetted.\n",
    "assert log_op.get_attr(\"fadd1\") is None\n",
    "\n",
    "# Check that fadd1 attr of other ops are intact.\n",
    "assert Op.get(\"exp\").get_attr(\"fadd1\")(1) == 2\n",
    "\n",
    "# Check that other attrs of the log op are intact.\n",
    "assert Op.get(\"log\").get_attr(\"fadd2\")(1) == 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9b21eef",
   "metadata": {},
   "source": [
    "## 缓存算子属性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1135df91",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "@dataclass\n",
    "class TempOpAttr:\n",
    "    \"\"\"Temporarily changes the attr of an op. Saves the required info for RAII pattern usage.\n",
    "    \n",
    "    Examples\n",
    "        --------\n",
    "        .. code-block:: python\n",
    "\n",
    "        # Temporarily update FTVMAlterOpLayout to a user-defined packed function.\n",
    "        # After the test is finished, the attr value will be set back to the original value.\n",
    "\n",
    "        with TempOpAttr(\"nn.conv2d\", \"FTVMAlterOpLayout\", alter_conv2d):\n",
    "            my_mod = relay.transform.AlterOpLayout()(my_mod)\n",
    "    \"\"\"\n",
    "    op_name : str # The op name.\n",
    "    attr_key : str # The attribute name.\n",
    "    attr_value : object # The attribute value.\n",
    "\n",
    "    def __post_init__(self):\n",
    "        self.op = Op.get(self.op_name)\n",
    "\n",
    "    def __enter__(self):\n",
    "        self.older_attr = self.op.get_attr(self.attr_key)\n",
    "        self.op.reset_attr(self.attr_key)\n",
    "        self.op.set_attr(self.attr_key, self.attr_value)\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, ptype, value, trace):\n",
    "        self.op.reset_attr(self.attr_key)\n",
    "        if self.older_attr:\n",
    "            self.op.set_attr(self.attr_key, self.older_attr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8246b30c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add1(x):\n",
    "    return x + 1\n",
    "\n",
    "def add2(x):\n",
    "    return x + 2\n",
    "\n",
    "# Set original attr value is add1.\n",
    "tvm.ir.register_op_attr(\"sqrt\", \"ftest\", add1)\n",
    "\n",
    "with TempOpAttr(\"sqrt\", \"ftest\", add2):\n",
    "    # Check that the attr value is updated to add2.\n",
    "    assert Op.get(\"sqrt\").get_attr(\"ftest\")(1) == 3\n",
    "\n",
    "# Check that the attr value is recovered to add1.\n",
    "assert Op.get(\"sqrt\").get_attr(\"ftest\")(1) == 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2ded086",
   "metadata": {},
   "source": [
    "## 算子级别"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bd18970c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm import relax\n",
    "x = relax.Var(\"x\")\n",
    "\n",
    "for op_name in [\"log\", \"exp\", \"sqrt\", \"rsqrt\", \"tanh\"]:\n",
    "    y = getattr(relax.op, op_name)(x)\n",
    "    assert y.op.name == f\"relax.{op_name}\"\n",
    "    assert y.op.support_level == 10\n",
    "    assert y.args[0] == x"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
